/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is dual-licensed under either the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree or the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree. You may select, at your option, one of the
 * above-listed licenses.
 */

use allocative::Allocative;
use indexmap::IndexSet;

use crate::analysis::registry::RecordedAnalysisValues;
use crate::artifact_groups::ArtifactGroup;

/// The result of evaluating a bxl function
#[derive(Allocative)]
pub struct BxlResult {
    /// The output string bytes generated by `ctx.output.print` in bxl
    output: Vec<u8>,
    /// The error string bytes
    error: Vec<u8>,
    /// The streaming output string bytes. This is where all streaming outputs not related to ensured artifacts will be written to.
    streaming: Vec<u8>,
    /// The artifacts that were called `ensure` during the execution of the bxl function.
    artifacts: Vec<ArtifactGroup>,
    /// The streaming output that is waiting on ensured artifacts to be materialized.
    pending_streaming_outputs: Vec<PendingStreamingOutput>,
    analysis_values: RecordedAnalysisValues,
}

impl BxlResult {
    pub fn new(
        output: Vec<u8>,
        error: Vec<u8>,
        streaming: Vec<u8>,
        ensured_artifacts: IndexSet<ArtifactGroup>,
        pending_streaming_outputs: Vec<(IndexSet<ArtifactGroup>, Vec<u8>)>,
        analysis_values: RecordedAnalysisValues,
    ) -> Self {
        Self {
            output,
            error,
            streaming,
            artifacts: ensured_artifacts.into_iter().collect(),
            pending_streaming_outputs: pending_streaming_outputs
                .into_iter()
                .map(|(waits_on, output)| PendingStreamingOutput::new(waits_on, output))
                .collect(),
            analysis_values,
        }
    }

    pub(crate) fn analysis_values(&self) -> &RecordedAnalysisValues {
        &self.analysis_values
    }

    pub fn output(&self) -> &Vec<u8> {
        &self.output
    }

    pub fn error(&self) -> &Vec<u8> {
        &self.error
    }

    pub fn streaming(&self) -> &Vec<u8> {
        &self.streaming
    }

    pub fn artifacts(&self) -> &Vec<ArtifactGroup> {
        &self.artifacts
    }

    pub fn pending_streaming_outputs(&self) -> &Vec<PendingStreamingOutput> {
        &self.pending_streaming_outputs
    }
}

#[derive(Allocative, Debug, Clone)]
pub struct PendingStreamingOutput {
    waits_on: IndexSet<ArtifactGroup>,
    output: Vec<u8>,
}

impl PendingStreamingOutput {
    pub fn new(waits_on: IndexSet<ArtifactGroup>, output: Vec<u8>) -> Self {
        Self { waits_on, output }
    }

    pub fn waits_on(&self) -> &IndexSet<ArtifactGroup> {
        &self.waits_on
    }

    pub fn output(&self) -> &[u8] {
        &self.output
    }

    pub fn remove_wait_on(&mut self, artifact: &ArtifactGroup) -> bool {
        self.waits_on.swap_remove(artifact)
    }

    pub fn is_pending(&self) -> bool {
        !self.waits_on.is_empty()
    }
}
